package com.luckzj.tools.agent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * MyAgent.java
 * 对log4j2 的JndiLookup进行修改，使其不再查找jndi对象
 * Author: zhoujie
 * CreateTime: 2021-12-10
 */
public class Log4jAgent {
    private static final String JNDI_LOOKUP_CLASS = "org/apache/logging/log4j/core/lookup/JndiLookup";

    public static void premain(String agentArgs, Instrumentation inst) {
        inst.addTransformer(new MyTransformer(), true);
    }

    static class MyTransformer implements ClassFileTransformer {

        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined
                , ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

            if (JNDI_LOOKUP_CLASS.equals(className)) {
                System.out.println("JNDILookUp Found!");

                try {
                    ClassPool classPool = ClassPool.getDefault();
                    CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
                    CtMethod method = ctClass.getDeclaredMethod("lookup");
                    method.setBody("return \"suspicious log msg: jndi:\" + $2;");

                    classfileBuffer = ctClass.toBytecode();

                    System.out.println("JNDILookUp Transformed!");

                } catch (Exception e) {
                    System.out.println("JNDILookUp Transform Error!" + e.getMessage());
                }
                finally {
                    System.out.println("finally");
                }
            }

            return classfileBuffer;
        }
    }
}
